import numpy as np
import pandas as pd

class CorrcoefCalculator():
    def __init__(self):
        self.Ex, self.Ey, self.Exy = 0, 0, 0
        self.Ex2, self.Ey2 = 0, 0
        self.traces_num = 0

    def update(self, traces, guess_midvalue):
        hyp=guess_midvalue.transpose()
        traces_num = traces.shape[0]
        self.Ex = (self.traces_num * self.Ex + np.sum(traces, axis=0)) / (self.traces_num + traces_num)
        self.Ey = (self.traces_num * self.Ey + np.sum(hyp, axis=1)) / (self.traces_num + traces_num)
        self.Exy = (self.traces_num * self.Exy + np.sum(traces * hyp[:, :, None], axis=1)) / (self.traces_num + traces_num)
        self.Ex2 = (self.traces_num * self.Ex2 + np.sum(traces ** 2, axis=0)) / (self.traces_num + traces_num)
        self.Ey2 = (self.traces_num * self.Ey2 + np.sum(hyp ** 2, axis=1)) / (self.traces_num + traces_num)
        self.traces_num+=traces_num

    def get_corrcoef(self):
        cov = self.Exy - self.Ey[:, None] * self.Ex[None, :]
        std_x = np.sqrt(self.Ex2 - self.Ex**2)
        std_y = np.sqrt(self.Ey2 - self.Ey**2)
        std_x = pd.DataFrame(std_x).replace(0, np.inf).to_numpy()[:, 0]
        std_y = pd.DataFrame(std_y).replace(0, np.inf).to_numpy()[:, 0]
        corrcoef = cov / (std_y[:, None] * std_x[None, :])
        return corrcoef

def attack(self):
    trs_num=self.traces_attack.shape[0]
    attack_steps=self.param_dict["attack_steps"]
    index = [i for i in range(0, trs_num, attack_steps)]
    index.append(trs_num)
    midvalue=self.midvalue_attack
    if len(midvalue.shape)==3:
        if self.target_bit is None:
            midvalue=midvalue@[2**(7-i) for i in range(8)]
        else:
            midvalue=midvalue[:,:,self.target_bit]
    corrcoef_calculator = CorrcoefCalculator()
    for i in range(len(index) - 1):
        corrcoef_calculator.update(self.traces_attack[index[i]:index[i+1]], midvalue[index[i]:index[i+1]])
        corrcoef=corrcoef_calculator.get_corrcoef()
        corrcoef=np.max(np.abs(corrcoef), axis=1)
        yield corrcoef,index[i+1]

def get_guess_key(self, corrcoef):
    guess_key = np.argmin(-corrcoef)
    return guess_key

def get_rank(self, corrcoef,real_key):
    corrcoef_sort=np.argsort(np.argsort(-corrcoef))
    return corrcoef_sort[real_key]